package org.example.cases.single;

import com.antgroup.geaflow.api.function.base.FilterFunction;
import com.antgroup.geaflow.api.function.base.MapFunction;
import com.antgroup.geaflow.api.graph.PGraphWindow;
import com.antgroup.geaflow.api.graph.compute.VertexCentricCompute;
import com.antgroup.geaflow.api.graph.function.vc.VertexCentricCombineFunction;
import com.antgroup.geaflow.api.graph.function.vc.VertexCentricComputeFunction;
import com.antgroup.geaflow.api.pdata.stream.window.PWindowSource;
import com.antgroup.geaflow.api.window.impl.AllWindow;
import com.antgroup.geaflow.common.config.Configuration;
import com.antgroup.geaflow.env.Environment;
import com.antgroup.geaflow.env.EnvironmentFactory;
import com.antgroup.geaflow.model.graph.edge.EdgeDirection;
import com.antgroup.geaflow.model.graph.edge.IEdge;
import com.antgroup.geaflow.model.graph.edge.impl.ValueEdge;
import com.antgroup.geaflow.model.graph.vertex.IVertex;
import com.antgroup.geaflow.model.graph.vertex.impl.ValueVertex;
import com.antgroup.geaflow.pipeline.IPipelineResult;
import com.antgroup.geaflow.pipeline.Pipeline;
import com.antgroup.geaflow.pipeline.PipelineFactory;
import com.antgroup.geaflow.pipeline.task.PipelineTask;
import com.antgroup.geaflow.view.GraphViewBuilder;
import com.antgroup.geaflow.view.IViewDesc.BackendType;
import com.antgroup.geaflow.view.graph.GraphViewDesc;
import org.example.config.JobConfig;
import org.example.function.AbstractVcFunc;
import org.example.function.CSVFileSink;
import org.example.function.CSVFileSinkV2;
import org.example.function.CSVFileSource;
import org.example.model.LongStringResult;
import org.example.util.FileUtils;

import java.util.Arrays;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;

public class Case3Single {

    private final String snapshotPath;

    private final String resultPath;

    public Case3Single(String snapshotPath, String resultPath) {
        this.snapshotPath = snapshotPath;
        this.resultPath = resultPath;
    }

    public static void main(String[] args) {
        long st = System.currentTimeMillis();
        new Case3Single(args[0], args[1]).start();
        long ed = System.currentTimeMillis();
        System.out.println("cost time: " + (ed - st) + " ms");
    }

    public void start() {
        Environment environment = EnvironmentFactory.onLocalEnvironment();
        IPipelineResult<?> result = submit(environment);
        result.get();
        environment.shutdown();
    }

    public IPipelineResult<?> submit(Environment environment) {
        Pipeline pipeline = PipelineFactory.buildPipeline(environment);
        Configuration envConfig = environment.getEnvironmentContext().getConfig();
        envConfig.put(CSVFileSink.OUTPUT_DIR, this.resultPath);

        pipeline.submit((PipelineTask) pipelineTaskCxt -> {
            Configuration conf = pipelineTaskCxt.getConfig();
            PWindowSource<IVertex<Long, Double>> vertices = pipelineTaskCxt.buildSource(new CSVFileSource<>(
                            FileUtils.pathJoin(this.snapshotPath, "Account.csv"),
                            Collections.singletonList("accountId"),
                            (CSVFileSource.CSVLineParser<IVertex<Long, Double>>) line -> {
                                IVertex<Long, Double> vertex = new ValueVertex<>(Long.valueOf(line.get(0)), 0.0);
                                return Collections.singletonList(vertex);
                            }
                    ), AllWindow.getInstance())
                    .withParallelism(conf.getInteger(JobConfig.SOURCE_PARALLELISM));

            PWindowSource<IEdge<Long, Double>> edges = pipelineTaskCxt.buildSource(new CSVFileSource<>(
                            FileUtils.pathJoin(this.snapshotPath, "AccountTransferAccount.csv"),
                            Arrays.asList("fromId", "toId", "amount"),
                            (CSVFileSource.CSVLineParser<IEdge<Long, Double>>) line -> {
                                Long src = Long.valueOf(line.get(0));
                                Long dst = Long.valueOf(line.get(1));
                                Double value = Double.valueOf(line.get(2));
                                IEdge<Long, Double> edge1 = new ValueEdge<>(src, dst, value, EdgeDirection.OUT);
                                IEdge<Long, Double> edge2 = new ValueEdge<>(dst, src, value, EdgeDirection.IN);
                                return Arrays.asList(edge1, edge2);
                            }
                    ), AllWindow.getInstance())
                    .withParallelism(conf.getInteger(JobConfig.SINK_PARALLELISM));

            int iterationParallelism = conf.getInteger(JobConfig.ITERATOR_PARALLELISM);
            GraphViewDesc graphViewDesc = GraphViewBuilder
                    .createGraphView(GraphViewBuilder.DEFAULT_GRAPH)
                    .withShardNum(iterationParallelism)
                    .withBackend(BackendType.Memory)
                    .build();

            PGraphWindow<Long, Double, Double> graphWindow =
                    pipelineTaskCxt.buildWindowStreamGraph(vertices, edges, graphViewDesc);

            graphWindow.compute(new PRAlgorithms(1))
                    .compute(iterationParallelism)
                    .getVertices()
                    .filter((FilterFunction<IVertex<Long, Double>>) record -> record.getValue() != 0)
                    .withParallelism(conf.getInteger(JobConfig.FILTER_PARALLELISM))
                    .map((MapFunction<IVertex<Long, Double>, LongStringResult>) value ->
                            new LongStringResult(value.getId(), String.format("%.2f", value.getValue())))
                    .withParallelism(conf.getInteger(JobConfig.MAP_PARALLELISM))
                    .sink(new CSVFileSinkV2<>("result3.csv"))
                    .withParallelism(conf.getInteger(JobConfig.SINK_PARALLELISM));

        });
        return pipeline.execute();
    }

    public static class PRAlgorithms extends VertexCentricCompute<Long, Double, Double, Double> {

        public PRAlgorithms(long iterations) {
            super(iterations);
        }

        @Override
        public VertexCentricComputeFunction<Long, Double, Double, Double> getComputeFunction() {
            return new PRVertexCentricComputeFunction();
        }

        @Override
        public VertexCentricCombineFunction<Double> getCombineFunction() {
            return null;
        }
    }

    public static class PRVertexCentricComputeFunction extends AbstractVcFunc<Long, Double, Double, Double> {
        @Override
        public void compute(Long vertexId, Iterator<Double> messageIterator) {
            List<IEdge<Long, Double>> outEdges = this.context.edges().getOutEdges();
            List<IEdge<Long, Double>> inEdges = this.context.edges().getInEdges();
            if (!outEdges.isEmpty() && !inEdges.isEmpty()) {
                double inSum = 0.0, outSum = 0.0;
                for (IEdge<Long, Double> outEdge : outEdges) {
                    outSum += outEdge.getValue();
                }
                for (IEdge<Long, Double> inEdge : inEdges) {
                    inSum += inEdge.getValue();
                }
                this.context.setNewVertexValue(inSum / outSum);
            }
        }
    }
}
